#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
from retrying import retry
import requests
import os
import json


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


@retry
def GetNuccoreID(GenbankID):
    url = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi'
    payload = {"db": "nuccore", "term": GenbankID, "retmode": "json"}
    r = requests.get(url, params=payload)
    if r.status_code == 200:
        result_json = r.json()
        if len(result_json["esearchresult"]["idlist"]) > 0:
            return result_json["esearchresult"]["idlist"][0]
        else:
            return None
    else:
        print("Web Error1. Retrying ...")
        raise IOError("Web Error.")


@retry
def GetNuccoreFile(GenbankInnerID, GenbankID):
    url = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi'
    payload = {"db": "nuccore", "id": GenbankInnerID, "rettype": "fasta"}
    r = requests.get(url, params=payload)
    if r.status_code == 200:
        with open("D:/16SrRNA/" + GenbankID + ".fasta", 'wb') as fd:
            for chunk in r.iter_content(1024 * 1024):
                fd.write(chunk)
            fd.close()
        print(GenbankID + ".fasta downloaded.")
    else:
        print("Web Error2. Retrying ...")
        raise IOError("Web Error.")


if __name__ == '__main__':
    client, collection = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")

    for result in collection.find({"Genbank_16SrRNA": {"$exists": 1}}, {"_id": 0, "Genbank_16SrRNA": 1}):
        GenbankID = result["Genbank_16SrRNA"]
        if os.path.exists("D:/16SrRNA/" + GenbankID + ".fasta"):
            continue
        else:
            GenbankInnerID = GetNuccoreID(GenbankID)
            if GenbankInnerID != None:
                GetNuccoreFile(GenbankInnerID, GenbankID)

    CloseMongoDB(client)
